use crate::{
    ast::Expression,
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{values::BasicValueEnum, AddressSpace, IntPredicate};

impl<'ctx> LLVMBackend<'ctx> {

    /// Generate string.replace(old, new) method - replaces all occurrences of old with new
    ///
    /// Takes two arguments:
    /// - old: The substring to find and replace
    /// - new: The substring to replace with
    ///
    /// Returns: New string with all occurrences replaced
    ///
    /// Algorithm:
    /// 1. Validate arguments (exactly 2)
    /// 2. Handle edge cases (empty old string, old not found)
    /// 3. Count occurrences using strstr()
    /// 4. Calculate new string length
    /// 5. Allocate and build new string
    /// 6. Return result
    pub(in crate::llvm_backend) fn generate_string_replace_method(
        &mut self,
        var_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        // Validate arguments
        if arguments.len() != 2 {
            return Err(CompilerError::internal(
                "replace() requires exactly two arguments (old: String, new: String)",
            ));
        }

        // Get the string pointer to replace
        let string_ptr_ptr = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| CompilerError::internal(&format!("Undefined variable: {}", var_name)))?
            .into_pointer_value();

        // Load the string pointer value (i8*)
        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let string_value = self
            .builder
            .build_load(i8_ptr_type, string_ptr_ptr, "string_val")?;
        let string_ptr = string_value.into_pointer_value();

        // Generate old and new string values
        let old_string = self.generate_expression(&arguments[0])?;
        let new_string = self.generate_expression(&arguments[1])?;

        if !old_string.is_pointer_value() || !new_string.is_pointer_value() {
            return Err(CompilerError::internal(
                "replace() arguments must be strings",
            ));
        }

        let old_str_ptr = old_string.into_pointer_value();
        let new_str_ptr = new_string.into_pointer_value();

        // Get LLVM types
        let i8_type = self.context.i8_type();
        let i8_ptr_type = i8_type.ptr_type(AddressSpace::default());
        let i32_type = self.context.i32_type();
        let i64_type = self.context.i64_type();

        // Declare C library functions
        // strlen function
        let strlen_fn_type = i32_type.fn_type(&[i8_ptr_type.into()], false);
        let strlen_fn = if let Some(func) = self.module.get_function("strlen") {
            func
        } else {
            self.module.add_function("strlen", strlen_fn_type, None)
        };

        // strstr function - finds substring in string
        let strstr_fn_type = i8_ptr_type.fn_type(&[i8_ptr_type.into(), i8_ptr_type.into()], false);
        let strstr_fn = if let Some(func) = self.module.get_function("strstr") {
            func
        } else {
            self.module.add_function("strstr", strstr_fn_type, None)
        };

        // malloc function
        let malloc_fn_type = i8_ptr_type.fn_type(&[i64_type.into()], false);
        let malloc_fn = if let Some(func) = self.module.get_function("malloc") {
            func
        } else {
            self.module.add_function("malloc", malloc_fn_type, None)
        };

        // memcpy function - copy memory
        let memcpy_fn_type = i8_ptr_type.fn_type(
            &[i8_ptr_type.into(), i8_ptr_type.into(), i64_type.into()],
            false,
        );
        let memcpy_fn = if let Some(func) = self.module.get_function("memcpy") {
            func
        } else {
            self.module.add_function("memcpy", memcpy_fn_type, None)
        };

        // Get lengths of strings
        let orig_len_call =
            self.builder
                .build_call(strlen_fn, &[string_ptr.into()], "orig_len_call")?;
        let orig_len = orig_len_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();

        let old_len_call =
            self.builder
                .build_call(strlen_fn, &[old_str_ptr.into()], "old_len_call")?;
        let old_len = old_len_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();

        let new_len_call =
            self.builder
                .build_call(strlen_fn, &[new_str_ptr.into()], "new_len_call")?;
        let new_len = new_len_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();

        // Check if old string is empty
        let zero_i32 = i32_type.const_int(0, false);
        let old_is_empty = self.builder.build_int_compare(
            IntPredicate::EQ,
            old_len,
            zero_i32,
            "old_is_empty",
        )?;

        // Create blocks for control flow
        let current_block = self.builder.get_insert_block().unwrap();
        let function = current_block.get_parent().unwrap();

        let empty_old_block = self.context.append_basic_block(function, "empty_old");
        let count_block = self
            .context
            .append_basic_block(function, "count_occurrences");
        let merge_block = self.context.append_basic_block(function, "merge");

        // Branch based on whether old string is empty
        self.builder
            .build_conditional_branch(old_is_empty, empty_old_block, count_block)?;

        // Handle empty old string case: return original string
        self.builder.position_at_end(empty_old_block);
        let empty_old_result = string_ptr;
        self.builder.build_unconditional_branch(merge_block)?;

        // Count occurrences of old string
        self.builder.position_at_end(count_block);

        // Create counter variable
        let count_ptr = self.builder.build_alloca(i32_type, "count_ptr")?;
        self.builder.build_store(count_ptr, zero_i32)?;

        // Create position tracker
        let pos_ptr = self.builder.build_alloca(i8_ptr_type, "pos_ptr")?;
        self.builder.build_store(pos_ptr, string_ptr)?;

        // Loop to count occurrences
        let count_loop_block = self.context.append_basic_block(function, "count_loop");
        let count_loop_body = self.context.append_basic_block(function, "count_loop_body");
        let count_loop_end = self.context.append_basic_block(function, "count_loop_end");

        self.builder.build_unconditional_branch(count_loop_block)?;
        self.builder.position_at_end(count_loop_block);

        // Find next occurrence
        let current_pos = self
            .builder
            .build_load(i8_ptr_type, pos_ptr, "current_pos")?
            .into_pointer_value();
        let strstr_call = self.builder.build_call(
            strstr_fn,
            &[current_pos.into(), old_str_ptr.into()],
            "strstr_call",
        )?;
        let found_ptr = strstr_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Check if found
        let null_ptr = i8_ptr_type.const_null();
        let found = self.builder.build_int_compare(
            IntPredicate::NE,
            found_ptr,
            null_ptr,
            "found",
        )?;

        self.builder
            .build_conditional_branch(found, count_loop_body, count_loop_end)?;

        // Loop body: increment counter and move position
        self.builder.position_at_end(count_loop_body);
        let current_count = self
            .builder
            .build_load(i32_type, count_ptr, "current_count")?
            .into_int_value();
        let one_i32 = i32_type.const_int(1, false);
        let new_count = self
            .builder
            .build_int_add(current_count, one_i32, "new_count")?;
        self.builder.build_store(count_ptr, new_count)?;

        // Move position past the found occurrence
        unsafe {
            let next_pos =
                self.builder
                    .build_in_bounds_gep(i8_type, found_ptr, &[old_len], "next_pos")?;
            self.builder.build_store(pos_ptr, next_pos)?;
        }

        self.builder.build_unconditional_branch(count_loop_block)?;

        // After counting loop
        self.builder.position_at_end(count_loop_end);
        let final_count = self
            .builder
            .build_load(i32_type, count_ptr, "final_count")?
            .into_int_value();

        // Check if count is zero (no occurrences found)
        let count_is_zero = self.builder.build_int_compare(
            IntPredicate::EQ,
            final_count,
            zero_i32,
            "count_is_zero",
        )?;

        let no_replace_block = self.context.append_basic_block(function, "no_replace");
        let do_replace_block = self.context.append_basic_block(function, "do_replace");

        self.builder
            .build_conditional_branch(count_is_zero, no_replace_block, do_replace_block)?;

        // No occurrences found: return original string
        self.builder.position_at_end(no_replace_block);
        let no_replace_result = string_ptr;
        self.builder.build_unconditional_branch(merge_block)?;

        // Do replacement
        self.builder.position_at_end(do_replace_block);

        // Calculate new string length: orig_len - (count * old_len) + (count * new_len)
        let count_times_old =
            self.builder
                .build_int_mul(final_count, old_len, "count_times_old")?;
        let count_times_new =
            self.builder
                .build_int_mul(final_count, new_len, "count_times_new")?;
        let len_after_remove =
            self.builder
                .build_int_sub(orig_len, count_times_old, "len_after_remove")?;
        let new_total_len =
            self.builder
                .build_int_add(len_after_remove, count_times_new, "new_total_len")?;

        // Allocate memory for new string (length + 1 for null terminator)
        let one_i32_const = i32_type.const_int(1, false);
        let alloc_size = self
            .builder
            .build_int_add(new_total_len, one_i32_const, "alloc_size")?;
        let alloc_size_i64 = self
            .builder
            .build_int_cast(alloc_size, i64_type, "alloc_size_i64")?;

        let malloc_call =
            self.builder
                .build_call(malloc_fn, &[alloc_size_i64.into()], "malloc_call")?;
        let result_ptr = malloc_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Build the new string
        let src_pos_ptr = self.builder.build_alloca(i8_ptr_type, "src_pos_ptr")?;
        let dst_pos_ptr = self.builder.build_alloca(i8_ptr_type, "dst_pos_ptr")?;
        let remaining_count_ptr = self.builder.build_alloca(i32_type, "remaining_count_ptr")?;

        self.builder.build_store(src_pos_ptr, string_ptr)?;
        self.builder.build_store(dst_pos_ptr, result_ptr)?;
        self.builder.build_store(remaining_count_ptr, final_count)?;

        // Loop to build new string
        let build_loop_block = self.context.append_basic_block(function, "build_loop");
        let build_loop_body = self.context.append_basic_block(function, "build_loop_body");
        let build_loop_end = self.context.append_basic_block(function, "build_loop_end");

        self.builder.build_unconditional_branch(build_loop_block)?;
        self.builder.position_at_end(build_loop_block);

        // Check if we have more occurrences to replace
        let remaining = self
            .builder
            .build_load(i32_type, remaining_count_ptr, "remaining")?
            .into_int_value();
        let has_more = self.builder.build_int_compare(
            IntPredicate::SGT,
            remaining,
            zero_i32,
            "has_more",
        )?;

        self.builder
            .build_conditional_branch(has_more, build_loop_body, build_loop_end)?;

        // Loop body
        self.builder.position_at_end(build_loop_body);

        let current_src = self
            .builder
            .build_load(i8_ptr_type, src_pos_ptr, "current_src")?
            .into_pointer_value();
        let current_dst = self
            .builder
            .build_load(i8_ptr_type, dst_pos_ptr, "current_dst")?
            .into_pointer_value();

        // Find next occurrence
        let next_occurrence_call = self.builder.build_call(
            strstr_fn,
            &[current_src.into(), old_str_ptr.into()],
            "next_occurrence_call",
        )?;
        let next_occurrence = next_occurrence_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Calculate distance to occurrence (in bytes)
        let src_int = self
            .builder
            .build_ptr_to_int(current_src, i64_type, "src_int")?;
        let occur_int = self
            .builder
            .build_ptr_to_int(next_occurrence, i64_type, "occur_int")?;
        let distance = self.builder.build_int_sub(occur_int, src_int, "distance")?;
        let distance_i32 = self
            .builder
            .build_int_cast(distance, i32_type, "distance_i32")?;

        // Copy text before the occurrence
        let distance_i64 = self
            .builder
            .build_int_cast(distance_i32, i64_type, "distance_i64")?;
        self.builder.build_call(
            memcpy_fn,
            &[current_dst.into(), current_src.into(), distance_i64.into()],
            "copy_before",
        )?;

        // Move dst pointer past copied text
        unsafe {
            let new_dst = self.builder.build_in_bounds_gep(
                i8_type,
                current_dst,
                &[distance_i32],
                "new_dst",
            )?;
            self.builder.build_store(dst_pos_ptr, new_dst)?;
        }

        // Copy new string
        let current_dst2 = self
            .builder
            .build_load(i8_ptr_type, dst_pos_ptr, "current_dst2")?
            .into_pointer_value();
        let new_len_i64 = self
            .builder
            .build_int_cast(new_len, i64_type, "new_len_i64")?;
        self.builder.build_call(
            memcpy_fn,
            &[current_dst2.into(), new_str_ptr.into(), new_len_i64.into()],
            "copy_new",
        )?;

        // Move dst pointer past new string
        unsafe {
            let new_dst2 =
                self.builder
                    .build_in_bounds_gep(i8_type, current_dst2, &[new_len], "new_dst2")?;
            self.builder.build_store(dst_pos_ptr, new_dst2)?;
        }

        // Move src pointer past old string
        unsafe {
            let new_src = self.builder.build_in_bounds_gep(
                i8_type,
                next_occurrence,
                &[old_len],
                "new_src",
            )?;
            self.builder.build_store(src_pos_ptr, new_src)?;
        }

        // Decrement remaining count
        let new_remaining = self
            .builder
            .build_int_sub(remaining, one_i32, "new_remaining")?;
        self.builder
            .build_store(remaining_count_ptr, new_remaining)?;

        self.builder.build_unconditional_branch(build_loop_block)?;

        // After build loop: copy remaining text
        self.builder.position_at_end(build_loop_end);

        let final_src = self
            .builder
            .build_load(i8_ptr_type, src_pos_ptr, "final_src")?
            .into_pointer_value();
        let final_dst = self
            .builder
            .build_load(i8_ptr_type, dst_pos_ptr, "final_dst")?
            .into_pointer_value();

        // Calculate remaining length
        let final_src_call =
            self.builder
                .build_call(strlen_fn, &[final_src.into()], "final_src_len_call")?;
        let final_src_len = final_src_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();
        let final_src_len_i64 =
            self.builder
                .build_int_cast(final_src_len, i64_type, "final_src_len_i64")?;

        // Copy remaining text
        self.builder.build_call(
            memcpy_fn,
            &[final_dst.into(), final_src.into(), final_src_len_i64.into()],
            "copy_remaining",
        )?;

        // Add null terminator
        unsafe {
            let null_pos = self.builder.build_in_bounds_gep(
                i8_type,
                final_dst,
                &[final_src_len],
                "null_pos",
            )?;
            let zero_i8 = i8_type.const_int(0, false);
            self.builder.build_store(null_pos, zero_i8)?;
        }

        let replace_result = result_ptr;
        let build_loop_end_block = self.builder.get_insert_block().unwrap();
        self.builder.build_unconditional_branch(merge_block)?;

        // Merge block: combine all results
        self.builder.position_at_end(merge_block);

        let result_phi = self.builder.build_phi(i8_ptr_type, "result_phi")?;
        result_phi.add_incoming(&[
            (&empty_old_result, empty_old_block),
            (&no_replace_result, no_replace_block),
            (&replace_result, build_loop_end_block),
        ]);

        Ok(result_phi.as_basic_value())
    }
}
